use std::path::PathBuf;

use anyhow::Error as E;
use candle_core::{DType, IndexOp, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::qwen2::{Config, ModelForCausalLM};
use clap::Parser;
use tokenizers::Tokenizer;

#[derive(Debug, Parser)]
#[command(version, about, long_about = None)]
struct Args {
    #[arg(short = 'p', long, default_value_t = String::from("/Users/c/.cache/huggingface/hub/models--Qwen--Qwen1.5-0.5B/snapshots/8f445e3628f3500ee69f24e1303c9f10f5342a39"))]
    model_path: String,
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();
    let path = PathBuf::from(args.model_path);
    let device = candle_core::Device::Cpu;
    let tokenizer = Tokenizer::from_file(path.join("tokenizer.json")).map_err(E::msg)?;
    let encoding = tokenizer.encode("你好", false).map_err(E::msg)?;
    let mut tokens = encoding.get_ids().to_vec();

    let vb = unsafe {
        VarBuilder::from_mmaped_safetensors(&[path.join("model.safetensors")], DType::F32, &device)?
    };

    let config: Config = serde_json::from_str(&std::fs::read_to_string(path.join("config.json"))?)?;

    let mut model = ModelForCausalLM::new(&config, vb)?;
    for i in 0..100 {
        let input = Tensor::new(tokens.as_slice(), &device)?.unsqueeze(0)?;
        println!("input tensor: {}", input);
        let output = model.forward(&input, i)?;
        println!("output tensor: {}", output);
        let next_token: u32 = output.squeeze(0)?.argmax(1)?.i(0)?.to_scalar()?;
        tokens.push(next_token);
        let res = tokenizer.decode(&[next_token], true).map_err(E::msg)?;
        println!("res: {}", res);
    }
    Ok(())
}
